package main

import (
	"flag"
	"fmt"
	"math/rand"
	"runtime"
	"strconv"
	"sync"
	"time"
)

// 定义日志函数
func logInfo(message string) {
	logTime := time.Now().Format("2006-01-02 15:04:05")
	fmt.Printf("[%s] info %s\n", logTime, message)
}

// 控制 CPU 占用的 Goroutine
func cpuConsumer(maxCores int, sleepDuration time.Duration) {
	if sleepDuration == 0 {
		sleepDuration = time.Duration(rand.Intn(5)+3) * time.Second
	}
	for {
		// 随机选择核心数量
		cores := rand.Intn(maxCores) + 1
		// 设置核心数量
		runtime.GOMAXPROCS(cores)
		// 打印当前占用的核心数量
		logInfo(fmt.Sprintf("Using %d cores", cores))
		// 增加工作负载
		for i := 0; i < 100000000; i++ {
		}
		// 随机等待一段时间
		time.Sleep(sleepDuration)
	}
}

// 控制内存占用的 Goroutine
func memoryConsumer(maxMemoryBytes int64, sleepDuration time.Duration) {
	if sleepDuration == 0 {
		sleepDuration = time.Duration(rand.Intn(5)+3) * time.Minute
	}
	for {
		// 随机选择内存占用
		memory := rand.Int63n(maxMemoryBytes)
		// 分配内存
		memoryBlock := make([]byte, memory)
		for i := range memoryBlock {
			memoryBlock[i] = byte(rand.Intn(256))
		}
		// 打印当前占用的内存
		logInfo(fmt.Sprintf("Using %d KB memory", memory/1024))
		// 随机等待一段时间
		time.Sleep(sleepDuration)
		// 释放内存
		memoryBlock = nil
		// 强制垃圾回收
		runtime.GC()
	}
}

func parseDuration(durationStr string) (time.Duration, error) {
	unit := durationStr[len(durationStr)-1]
	value, err := strconv.Atoi(durationStr[:len(durationStr)-1])
	if err != nil {
		return 0, err
	}

	switch unit {
	case 's', 'S':
		return time.Duration(value) * time.Second, nil
	case 'm', 'M':
		return time.Duration(value) * time.Minute, nil
	case 'h', 'H':
		return time.Duration(value) * time.Hour, nil
	default:
		return 0, fmt.Errorf("invalid duration format")
	}
}

func main() {
	// 定义命令行参数
	var maxCores int
	var maxMemory string
	var cpuSleepDurationStr string
	var memSleepDurationStr string

	flag.IntVar(&maxCores, "c", 1, "Maximum number of cores to use")
	flag.StringVar(&maxMemory, "m", "1G", "Maximum memory to use (e.g., 1G, 512M)")
	flag.StringVar(&cpuSleepDurationStr, "ctime", "", "CPU consumer sleep duration (e.g., s, m, h)")
	flag.StringVar(&memSleepDurationStr, "mtime", "", "Memory consumer sleep duration (e.g., s, m, h)")

	// 自定义帮助信息
	flag.Usage = func() {
		_, err := fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s:\n", "resource_consumer")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "  -c int\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "    	Maximum number of cores to use (default 1)\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "  -m string\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "    	Maximum memory to use (e.g., 1G, 512M) (default \"1G\")\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "  -ctime string\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "    	CPU consumer sleep duration (e.g., 5s, 5m, 5h)\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "  -mtime string\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
		_, err = fmt.Fprintf(flag.CommandLine.Output(), "    	Memory consumer sleep duration (e.g., 5s, 5m, 5h)\n")
		if err != nil {
			logInfo(fmt.Sprintf("Error writing usage information: %v", err))
		}
	}

	flag.Parse()

	// 解析最大内存
	var maxMemoryBytes int64
	switch maxMemory[len(maxMemory)-1] {
	case 'G', 'g':
		memValue, err := strconv.Atoi(maxMemory[:len(maxMemory)-1])
		if err != nil {
			logInfo(fmt.Sprintf("Invalid memory format: %v", err))
			return
		}
		maxMemoryBytes = int64(memValue) * 1024 * 1024 * 1024
	case 'M', 'm':
		memValue, err := strconv.Atoi(maxMemory[:len(maxMemory)-1])
		if err != nil {
			logInfo(fmt.Sprintf("Invalid memory format: %v", err))
			return
		}
		maxMemoryBytes = int64(memValue) * 1024 * 1024
	default:
		logInfo("Invalid memory format. Use G or M.")
		return
	}

	// 解析睡眠时间
	var cpuSleepDuration, memSleepDuration time.Duration
	var err error

	if cpuSleepDurationStr != "" {
		cpuSleepDuration, err = parseDuration(cpuSleepDurationStr)
		if err != nil {
			logInfo(fmt.Sprintf("Invalid CPU sleep duration format: %v", err))
			return
		}
	}

	if memSleepDurationStr != "" {
		memSleepDuration, err = parseDuration(memSleepDurationStr)
		if err != nil {
			logInfo(fmt.Sprintf("Invalid memory sleep duration format: %v", err))
			return
		}
	}

	// 设置最大核心数量
	runtime.GOMAXPROCS(maxCores)

	logInfo(fmt.Sprintf("Using up to %d cores and %s memory", maxCores, maxMemory))
	if cpuSleepDurationStr != "" {
		logInfo(fmt.Sprintf("CPU consumer sleep duration: %s", cpuSleepDuration))
	}
	if memSleepDurationStr != "" {
		logInfo(fmt.Sprintf("Memory consumer sleep duration: %s", memSleepDuration))
	}

	// 使用 WaitGroup 等待 Goroutine 结束
	var wg sync.WaitGroup
	wg.Add(2)

	// 启动 CPU 占用的 Goroutine
	go func() {
		defer wg.Done()
		cpuConsumer(maxCores, cpuSleepDuration)
	}()

	// 启动内存占用的 Goroutine
	go func() {
		defer wg.Done()
		memoryConsumer(maxMemoryBytes, memSleepDuration)
	}()

	// 等待所有 Goroutine 结束
	wg.Wait()
}
